# 高斯混合模型应用示例
import argparse
from typing import Dict
import numpy as np
from scipy.stats import norm

class GmmApp(object):
    def __init__(self):
        self.name = 'apps.tlp.gmm_app.GmmApp'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'雷电预测混合高斯模型 v0.0.3')
        # 设置随机种子以确保结果可重复
        np.random.seed(0)
        # 生成一个非常大的数据集（模拟）
        n_samples = 1000000  # 100万条数据
        mu1, sigma1 = 5, 1  # 第一个高斯分布的均值和标准差
        mu2, sigma2 = 10, 2  # 第二个高斯分布的均值和标准差
        weights = [0.4, 0.6]  # 两个高斯分布的混合比例
        # 生成数据
        data = np.concatenate([
            np.random.normal(mu1, sigma1, int(n_samples * weights[0])),
            np.random.normal(mu2, sigma2, int(n_samples * weights[1]))
        ])
        # 打乱数据
        np.random.shuffle(data)
        # 初始化模型参数
        n_components = 2  # 假设有两个高斯分布
        means = np.array([4.0, 11.0])  # 初始均值
        variances = np.array([1.0, 4.0])  # 初始方差
        weights = np.array([0.5, 0.5])  # 初始混合比例
        # 定义小批量大小
        batch_size = 1000
        # 定义学习率（用于逐步更新参数）
        learning_rate = 0.01
        # 小批量 EM 算法
        for i in range(0, n_samples, batch_size):
            # 获取当前小批量数据
            batch = data[i:i + batch_size]
            # E 步：计算责任值
            responsibilities = np.zeros((len(batch), n_components))
            for k in range(n_components):
                responsibilities[:, k] = weights[k] * norm.pdf(batch, means[k], np.sqrt(variances[k]))
            responsibilities /= responsibilities.sum(axis=1, keepdims=True)  # 归一化
            # M 步：更新参数
            Nk = responsibilities.sum(axis=0)  # 每个高斯分布的有效样本数
            weights = (1 - learning_rate) * weights + learning_rate * (Nk / len(batch))  # 更新混合比例
            for k in range(n_components):
                means[k] = (1 - learning_rate) * means[k] + learning_rate * (responsibilities[:, k] @ batch) / Nk[k]  # 更新均值
                variances[k] = (1 - learning_rate) * variances[k] + learning_rate * (responsibilities[:, k] @ (batch - means[k])**2) / Nk[k]  # 更新方差
            # 打印进度
            if i % 100000 == 0:
                print(f"已处理 {i} 条数据")
                print("当前均值:", means)
                print("当前方差:", variances)
                print("当前混合比例:", weights)
        # 输出最终参数
        print("最终均值:", means)
        print("最终方差:", variances)
        print("最终混合比例:", weights)
        # 给定一个数
        x = 7.5
        # 计算该数属于每个高斯分布的概率
        probs = np.array([weights[k] * norm.pdf(x, means[k], np.sqrt(variances[k])) for k in range(n_components)])
        probs /= probs.sum()  # 归一化
        print(f"数据点 {x} 属于每个高斯分布的概率: {probs}")


def main(params:Dict = {}) -> None:
    GmmApp.startup(params=params)

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--run_mode', action='store',
        type=int, default=1, dest='run_mode',
        help='run mode'
    )
    return parser.parse_args()

if '__main__' == __name__:
    args = parse_args()
    params = vars(args)
    main(params=params)